# Created by Gao Song on 2022-03-29

import pandas as pd
import numpy as np
from .rq import rq
from .trade_date import get_prev_trading_date
from datetime import datetime, timedelta


def get_current_price_stock(symbols):
    if len(symbols) == 1:
        return pd.Series({symbols[0]: rq.current_snapshot(symbols).last})
    return pd.Series([x.last for x in rq.current_snapshot(symbols)], index=symbols)


def get_current_pct_stock(symbols):
    now = datetime.now()
    td = timedelta(hours=5)

    if len(symbols) == 1:
        tick = rq.current_snapshot(symbols)
        return pd.Series({symbols[0]: tick.last / tick.prev_close - 1 if now - tick.datetime < td else 0})
    return pd.Series([tick.last / tick.prev_close - 1 if now - tick.datetime < td else 0 for tick in rq.current_snapshot(symbols)], index=symbols)


def get_all_stocks():
    return rq.all_instruments(type='CS').order_book_id


def get_ex_factor(symbols, d):
    ex = rq.get_ex_factor(symbols, start_date=d, end_date=d)
    ex.reset_index(inplace=True)
    ex.set_index('order_book_id', inplace=True)
    return ex['ex_factor']


def get_industry(symbols, name=False):
    if name:
        return rq.get_instrument_industry(symbols, source='citics_2019', level=1, date=None, market='cn')['first_industry_name']
    else:
        return rq.get_instrument_industry(symbols, source='citics_2019', level=1, date=None, market='cn')['first_industry_code']


def get_stock_details(symbols, field):
    res = rq.instruments(symbols, market='cn')
    return pd.Series([getattr(r, field) for r in res], index=symbols)


def get_paused_mask_2d(d):
    prev_d = get_prev_trading_date(d)
    symbols = get_all_stocks()
    res = rq.is_suspended(symbols, start_date=prev_d, end_date=d).any()
    return res

def get_close_r_stock(symbols, start_date, end_date):
    price = rq.get_price(
        symbols, start_date=start_date, end_date=end_date, fields='close', adjust_type='post'
    )['close'].unstack().T
    return price.apply(np.log).diff()
